#include <mm/kheap.h>
#include <mm/kmalloc.h>
#include <mm/paging.h>
#include <utils/string.h>
#include <utils/debug.h>

struct kheap *kheap;

struct page_directory *kernel_directory;

extern uint32 end_mark;
uint32 placement_address = (uint32)&end_mark;

void malloc_print() {
    if (kheap != 0) kheap_print(kheap);
}

void* __alloc(uint32 size, void **phys, uint8 align) {
    void *addr = 0;
    void *physical = 0;
    if (kheap != 0) {
        // heap is active, let's use it
#ifdef DEBUG_MALLOC
        kprintf("HEAP IS ACTIVE\n");
#endif
        addr = kheap_alloc(kheap, size, align);
        // get physical address
        struct page *page = get_page(kernel_directory, (uint32) addr);
        physical = (void*)(page->frame * PAGE_SIZE + ((uint32)addr & 0xFFF));
    } else {
        // heap is not active, let's just use placement allocation
        if (align) {
            ENSURE_PAGE_ALIGNED(placement_address);
        }
        addr = (void*)placement_address;
        physical = addr;
        placement_address += size;
    }
    if (phys != 0) {
        *phys = physical;
    }
    memset(addr, 0, size);
    return addr;
}

void kfree(void* addr) {
    if (kheap != 0) {
        kheap_free(kheap, addr);
        malloc_print();
    }
}

void malloc_init(uint32 mem_end) {
    uint32 i = 0; // for loop iterator
    
    paging_init(mem_end);
    
    kernel_directory = (struct page_directory*)kmalloca(sizeof(struct page_directory));
    kernel_directory->physical_address = (uint32)kernel_directory->tables_addresses;

    for (i = KHEAP_START;i < KHEAP_START+KHEAP_INITIAL_SIZE; i += PAGE_SIZE) 
        if (!get_page(kernel_directory, i)) 
            new_page(kernel_directory, i, 0x7);
       
 
    // keep enough space in memory for heap and used utility structures so far.
    // After heap is active we will use that instead of placement.
    i = 0;
    while(i  < placement_address+sizeof(struct kheap)) {
        struct page *p = get_page(kernel_directory, i);
        if (!p){
            p = new_page(kernel_directory, i, 0x7);
        }
        alloc_frame(p, 0);
        i+= PAGE_SIZE;
    }
  
    for (i = KHEAP_START;i < KHEAP_START+KHEAP_INITIAL_SIZE; i += PAGE_SIZE) 
        alloc_frame(get_page(kernel_directory, i), 0);
   

    set_page_directory(kernel_directory);
    
    kheap = kheap_create(KHEAP_START, KHEAP_START + KHEAP_INITIAL_SIZE, 0xCFFFF000, 0);
}
